package com.soushangxing.comfyuibackendv2.service

import com.soushangxing.comfyuibackendv2.bean.BackendQueueInfo
import com.soushangxing.comfyuibackendv2.bean.GlobalQueue
import com.soushangxing.comfyuibackendv2.bean.ResponseType
import com.soushangxing.comfyuibackendv2.entity.OrderEntity
import com.soushangxing.comfyuibackendv2.entity.OrderErrorEntity
import com.soushangxing.comfyuibackendv2.repository.OrderErrorRepository
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.beans.factory.annotation.Value
import org.springframework.stereotype.Service
import java.time.LocalDateTime

@Service
class GlobalQueueService {

    @Autowired
    lateinit var comfyUIService: ComfyUIService

    @Autowired
    lateinit var orderErrorRepository: OrderErrorRepository

    @Autowired
    lateinit var globalQueue: GlobalQueue

    @Value("\${queue.maxWait}")
    private lateinit var maxWait: Number

    @Value("\${queue.maxRetry}")
    private lateinit var maxRetry: Number

    /**
     * 新版本：后端管理队列。
     * 将队伍加入后端全局变量
     */
    fun backendQueue(taskId:String, userId:String, prompt:String) : ResponseType {
        try {
            // 检查当前队列数量
            if (globalQueue.globalBackendQueue.size >= maxWait.toInt()) {
                return ResponseType("队列已满，无法添加新任务", 1, "队列已满")
            }
            // 检查当前可用服务器数量
            if (globalQueue.globalServer.count { it.status == "运行中" } == 0) {
                return ResponseType("当前尚无可用服务器", 1, "当前尚无可用服务器")
            }

            val queueInfo = BackendQueueInfo()
            queueInfo.taskId = taskId
            queueInfo.userId = userId
            queueInfo.prompt = prompt
            queueInfo.status = "排队中"
            queueInfo.createTime = LocalDateTime.now()

            globalQueue.globalBackendQueue.add(queueInfo)


//          每次添加任务，是检查队列的时间点1（时间点2是每次任务完成,时间点3是服务器启动完成）
            checkQueue()

            return ResponseType(queueInfo,0,"任务排队完成")
        }catch (e:Exception){
            return ResponseType("任务排队失败:"+e.message)
        }
    }

    /**
     * 通用：指派任务开始
     * 让一个任务开始！
     */
    fun queuePromptStart(taskId: String, serverUrl:String) {
        // 遍历 globalBackendQueue 查找匹配的任务
        val task = globalQueue.globalBackendQueue.find { it.taskId == taskId }
//            查询Socket的客户端ID
        val server = globalQueue.globalServer.find { it.url == serverUrl }

        if (task != null && server != null) {
            println("尝试开始$taskId 分配给$serverUrl");

//             调用 queuePromptNew 函数
            val response = comfyUIService.queuePromptNew(task.taskId, task.userId, task.prompt, serverUrl, "client-"+server.id)

            // 检查响应的 code 是否为 0
            if (response.code == 0) {
                if (response.data is OrderEntity) {
                    val orderEntity = response.data as OrderEntity
                    task.promptId = orderEntity.promptId
                } else {
                    task.status = "任务下发失败"

                    val errorEntity = OrderErrorEntity()
                    errorEntity.createTime = LocalDateTime.now()
                    errorEntity.status = task.status
                    errorEntity.taskId = task.taskId
                    errorEntity.userId = task.userId
                    errorEntity.message = "任务下发后，返回结构异常！！"
                    orderErrorRepository.save(errorEntity)
                    globalQueue.globalBackendQueue.removeIf { it.taskId == taskId }
                }

            } else {
                // 任务下发失败，可能是服务器问题，在服务器配置的最大重试限度内，重试
                task.retry = task.retry + 1
                if (task.retry <= maxRetry.toInt()){
                    task.status = "排队中"
                    task.serverUrl = ""
                    checkQueue()
                } else {
                    task.status = "任务下发失败"

                    val errorEntity = OrderErrorEntity()
                    errorEntity.createTime = LocalDateTime.now()
                    errorEntity.status = task.status
                    errorEntity.taskId = task.taskId
                    errorEntity.userId = task.userId
                    errorEntity.message = "任务下发过程报错"
                    orderErrorRepository.save(errorEntity)
                    globalQueue.globalBackendQueue.removeIf { it.taskId == taskId }
                }
            }
        }
    }

    /**
     * 通用：检查队伍状态
     * 如果当前正在运行的任务数量少于最大允许的数量，就指派相应的任务去空余的服务器运行
     * 调用时机：1.添加任务、2.任务完成、3.服务启动
     */
    fun checkQueue(){
//        1. 获取当前全局任务队列中，正在进行的任务数 以及 当前运行中的服务器的总放行数
        val runningTasksCount = globalQueue.globalBackendQueue.count { it.status != "排队中" }
        val totalLetGo = globalQueue.globalServer
            .filter { it.status == "运行中" }
            .sumOf { it.taskLimit }

//        2. 如果任务数小于letGo，则尝试按照创建时间，先进先出指派最早排队的任务，使任务状态!=排队中的任务数量等于letGo
//        （如果排队中的任务不足就算了）
        if (runningTasksCount < totalLetGo) {
            val tasksToStart = totalLetGo - runningTasksCount

            val queuedTasks = globalQueue.globalBackendQueue
                .filter { it.status == "排队中" }
                .sortedBy { it.createTime }
                .take(tasksToStart)

            for (task in queuedTasks) {
                val availableServerUrl = getAvailableServer()
                if (availableServerUrl != null) {
                    queuePromptStart(task.taskId, availableServerUrl)
                    // 更新任务的 serverUrl
                    task.serverUrl = availableServerUrl
                } else {
                    // 如果没有可用的服务器，跳出循环
                    break
                }
            }
        }
    }


    /**
     * 获取一个可用的服务器的URL
     */
    fun getAvailableServer(): String? {
        val runningServers = globalQueue.globalServer.filter { it.status == "运行中" }
        val serverTaskCounts = runningServers.associateWith { server ->
            globalQueue.globalBackendQueue.count { it.serverUrl == server.url && it.status != "排队中" }
        }

        val availableServers = serverTaskCounts.entries
            .filter { (server, taskCount) -> taskCount < server.taskLimit }
            .map { it.key }

        return if (availableServers.isNotEmpty()) {
            availableServers.random().url
        } else {
            null
        }
    }
}
